#include <arpa/inet.h>
#include <errno.h>
#include <net/ethernet.h>
#include <net/if.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <netinet/udp.h>
#include <pcap/pcap.h>
#include <pcap/sll.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <string>
#include <unistd.h>

#define NANOSECONDS_PER_SECOND 1000000000L

// 版本号
static const char *_VERSION_ = "1.1.0";

typedef struct _input_conf_ {
    int ifindex_;
    long speed_;
    int repeat_;
    int port_;
    int count_;
    int fd_;
    std::string host_;
    sockaddr_in addr_;

    _input_conf_() {
        ifindex_ = 0;
        speed_ = 0;
        repeat_ = 1;
        port_ = 53;
        count_ = 0;
        fd_ = -1;
        host_ = "127.0.0.1";
    }

    /**
     * @brief 打开发送端口
     *
     */
    bool open() {
        fd_ = socket(AF_INET, SOCK_DGRAM, 0);
        if (fd_ == -1) {
            printf("socket error=%s\n", strerror(errno));
            return false;
        }

        if (ifindex_ != 0) {
            ip_mreqn mreqn;
            memset(&mreqn, 0, sizeof(mreqn));
            mreqn.imr_ifindex = ifindex_;
            if (setsockopt(fd_, IPPROTO_IP, IP_MULTICAST_IF, &mreqn,
                           sizeof(mreqn)) == -1) {
                printf("setsockopt error=%s\n", strerror(errno));
                return false;
            }
        }

        memset(&addr_, 0, sizeof(addr_));
        addr_.sin_family = AF_INET;
        addr_.sin_port = htons(port_);
        if (inet_pton(AF_INET, host_.c_str(), &addr_.sin_addr) <= 0) {
            printf("inet_pton failed of %s\n", host_.c_str());
            return false;
        }
        return true;
    }

    /**
     * @brief 发送数据
     *
     */
    bool send(const u_char *d, ssize_t len) {
        ssize_t n = sendto(fd_, d, len, 0, reinterpret_cast<sockaddr *>(&addr_),
                           sizeof(addr_));
        if (n != len) {
            printf("failed sendto error=%s\n", strerror(errno));
            return false;
        }
        return true;
    }
} input_conf;

/**
 * @brief 显示帮助信息
 *
 */
void usage() {
    printf(
        "udpreplay %s haibin.wang https://gitee.com/fly542/udpreplay\n"
        "usage: udpreplay [-i iface] [-h host] [-p port] [-s speed] [-r "
        "repeat] pcap\n"
        "\n"
        "  -v          版本号\n"
        "  -i iface    发送数据包的网卡名称\n"
        "  -h host     接收数据包的服务器地址, (默认 127.0.0.1)\n"
        "  -p port     接收数据包的服务器端口号, (默认 53)\n"
        "  -c count    成功发送的数据包个数 (>0时 到达计数则停止发送, 默认0)\n"
        "  -r repeat   使用pcap包循环发送次数 (-1 无限循环, 默认1)\n"
        "  -s speed    每秒发包个数需>=0, (0表示不限速, 默认 0)\n",
        _VERSION_);
}

/**
 * @brief 解析运行参数
 *
 */
bool parse_args(input_conf &conf, int argc, char *argv[]) {
    int opt;
    while ((opt = getopt(argc, argv, "i:r:t:h:p:s:c:v")) != -1) {
        switch (opt) {
            case 'v':
                printf("version : %s\n", _VERSION_);
                return false;
            case 'i':
                conf.ifindex_ = if_nametoindex(optarg);
                if (conf.ifindex_ == 0) {
                    printf("if_nametoindex error=%s\n", strerror(errno));
                    return false;
                }
                break;
            case 'p':
                conf.port_ = atoi(optarg);
                break;
            case 'h':
                conf.host_ = optarg;
                break;
            case 's':
                conf.speed_ = atol(optarg);
                if (conf.speed_ < 0) {
                    printf("speed must >=0\n");
                    return false;
                }
                break;
            case 'c':
                conf.count_ = atol(optarg);
                if (conf.count_ <= 0) {
                    printf("count must >0\n");
                    return false;
                }
                break;
            case 'r':
                conf.repeat_ = atoi(optarg);
                if (conf.repeat_ != -1 && conf.repeat_ <= 0) {
                    printf("repeat must be positive integer or -1\n");
                    return false;
                }
                break;
            default:
                usage();
                return false;
        }
    }
    if (optind >= argc) {
        usage();
        return false;
    }
    return true;
}

/**
 * @brief 处理DNS数据包
 *
 */
void process_dns_packet(const u_char *p, size_t len, input_conf &conf,
                        const long &total, long &success) {
    // 检查包长度是否足够
    if (len < sizeof(udphdr)) {
        printf("Packet too small to contain DNS data, seq=%ld\n", total);
        return;
    }

    // 找到 DNS 数据包位置 (此处假设 UDP 包已确认有效)
    const u_char *dns_data = p;
    if (!conf.send(dns_data, len)) {
        printf("Failed to send DNS packet, seq=%ld\n", total);
    } else {
        success++;
    }
}

// 发送ip协议数据
bool send_ip_data(const u_char *udp_start, const struct ip *package_ip,
                  input_conf &conf, long &total, long &success) {
    const udphdr *udp = reinterpret_cast<const udphdr *>(udp_start);
    if (package_ip->ip_p == IPPROTO_UDP) {
        size_t len = ntohs(udp->len) - sizeof(udphdr);
        process_dns_packet(udp_start + sizeof(udphdr), len, conf, total,
                           success);
    } else if (package_ip->ip_p == IPPROTO_TCP) {
        const tcphdr *tcp = reinterpret_cast<const tcphdr *>(udp_start);
        const u_char *tcp_data = reinterpret_cast<const u_char *>(tcp + 1);
        size_t len = ntohs(tcp->doff) * 4;
        process_dns_packet(tcp_data, len, conf, total, success);
    } else {
        return false;
    }
    return true;
}

// 处理链路层SLL类型
bool linktype_sll_deal(const pcap_pkthdr &header, const u_char *p,
                       input_conf &conf, long &total, long &success) {
    const struct sll_header *hdr =
        reinterpret_cast<const struct sll_header *>(p);
    const u_char *payload = p + sizeof(struct sll_header);
    unsigned short protocol = ntohs(hdr->sll_protocol);

    if (protocol == ETH_P_8021Q) {  // 0x8100 是VLAN头，需要跳过4字节VLAN标签
        if (header.caplen < sizeof(struct sll_header) + 4) {
            return false;
            // printf("Packet too small to contain VLAN data.seq=%ld\n", total);
        }
        payload += 4;
        protocol = ntohs(*reinterpret_cast<const u_int16_t *>(payload - 2));
    }

    if (protocol == ETH_P_IP) {  // 0x0800
        const struct ip *package_ip =
            reinterpret_cast<const struct ip *>(payload);
        if (package_ip->ip_v != 4) {
            return false;
            // printf("type=%d\n", package_ip->ip_v);
        }
        const u_char *udp_start = payload + package_ip->ip_hl * 4;
        return send_ip_data(udp_start, package_ip, conf, total, success);
    } else {
        return false;
        // printf("unknow type=0x%X", protocol);
    }
    return true;
}

// 处理链路层 EN10MB 类型
bool linktype_en10mb_deal(const u_char *p, input_conf &conf, long &total,
                          long &success) {
    const ether_header *eth = reinterpret_cast<const ether_header *>(p);
    if (ntohs(eth->ether_type) != ETHERTYPE_IP) {
        printf("Error: Ethernet type is %d\n", ntohs(eth->ether_type));
        return false;
    }
    const struct ip *package_ip =
        reinterpret_cast<const struct ip *>(p + sizeof(ether_header));
    if (package_ip->ip_v != 4) {
        return false;
    }
    const u_char *udp_start = p + sizeof(ether_header) + package_ip->ip_hl * 4;
    return send_ip_data(udp_start, package_ip, conf, total, success);
}

int main(int argc, char *argv[]) {
    input_conf conf;
    if (!parse_args(conf, argc, argv)) {
        return 1;
    }

    if (!conf.open()) {
        return 1;
    }
    int speed_nano_sleep = 0;
    if (conf.speed_ > 0) {
        speed_nano_sleep = NANOSECONDS_PER_SECOND / conf.speed_;
    }
    long success = 0;
    long total = 0;
    for (int i = 0; conf.repeat_ == -1 || i < conf.repeat_; i++) {
        bool bcount = false;
        char errbuf[PCAP_ERRBUF_SIZE];
        pcap_t *handle = pcap_open_offline_with_tstamp_precision(
            argv[optind], PCAP_TSTAMP_PRECISION_NANO, errbuf);

        if (NULL == handle) {
            printf("failed pcap_open error=%s\n", errbuf);
            return 1;
        }
        int linktype = pcap_datalink(handle);  // 数据链路层类型

        pcap_pkthdr header;
        const u_char *p;
        struct timeval start_time, end_time;
        gettimeofday(&start_time, NULL);
        while ((p = pcap_next(handle, &header))) {
            total++;
            if (DLT_LINUX_SLL == linktype) {  // DLT_LINUX_SLL=113
                if (!linktype_sll_deal(header, p, conf, total, success)) {
                    continue;
                }
            } else if (DLT_EN10MB == linktype) {  // DLT_EN10MB=1
                if (!linktype_en10mb_deal(p, conf, total, success)) {
                    continue;
                }
            } else {
                // Add handling for other link types if needed
                printf("Unsupported linktype: %d\n", linktype);
                // for (int i = 0; i < 64 && i < header.caplen; i++) {
                //    printf(" %02x", p[i]);
                //}
                // printf("\n");

                continue;
            }

            if (speed_nano_sleep > 0) {
                long long tt = speed_nano_sleep * success / 1000;
                gettimeofday(&end_time, NULL);
                long long usec =
                    (end_time.tv_sec - start_time.tv_sec) * 1000000 +
                    (end_time.tv_usec - start_time.tv_usec);
                int usl = tt - usec;
                if (usl > 0) {  // >1微秒则休息1微秒
                    usleep(usl);
                }
            }

            if (conf.count_ > 0 && success >= conf.count_) {
                bcount = true;
                break;
            }
        }
        gettimeofday(&end_time, NULL);
        pcap_close(handle);
        printf("total packets: %ld\nsuccess send udp packages:%ld\n", total,
               success);
        long long usec = (end_time.tv_sec - start_time.tv_sec) * 1000000 +
                         (end_time.tv_usec - start_time.tv_usec);
        printf("time cost:%lldms\nsend speed: %lld/s\n", usec / 1000,
               (success * 1000000 / usec));
        if (bcount) {
            break;
        }
    }

    return 0;
}
